Co-authored-by: Copilot <copilot@github.com>

Remove the `_ep_debug` function and all its calls throughout the expert parallel module to clean up production code. The debug logging was gated behind the `TWINKLE_EP_DEBUG` environment variable but added unnecessary complexity and performance overhead in normal operation. Also remove unused imports (`os`, `time`) and the `_ep_debug_name` attribute assignment on blocks.

The sync_after_backward feature and its associated configuration have been removed as they are no longer needed for maintaining EP/FSDP collective ordering. This simplifies the codebase by eliminating unused synchronization logic.

…ECONDS

Add support for customizing the distributed process group timeout through the environment variable `TWINKLE_DIST_TIMEOUT_SECONDS`, defaulting to 7200 seconds (2 hours). This prevents timeout errors in long-running distributed training jobs.

…ffers

Some Transformers models keep persistent buffers in state_dict that are not DTensors. The original accelerate FSDP2 load function assumed all entries have device_mesh, causing failures. This patch adds a monkey-patch to handle both DTensor parameters and regular tensor buffers during state dict loading.

When loading full state dict in FSDP2, the parameter might not exist in the model (e.g., when using tied weights or shared parameters). This change adds a try-except to handle AttributeError when getting the parameter, and only applies contiguous conversion when the parameter exists and is contiguous.

…e FSDP configuration

Add helper functions for parsing boolean, torch dtype, and lora target modules from environment variables. Introduce new config options for low CPU memory usage, memory efficient init, model dtype, and flexible FSDP/DP/EP sizing. Improve accelerate FSDP2 state dict loading with parameter dtype inference and contiguous casting.

Add debug logging to FSDP2 state dict loading and model preparation, controlled by TWINKLE_FSDP_DEBUG environment variable. Also optimize by returning early with original implementation when running on CUDA devices.

…d state loading traces

- Replace logger.info with print in barrier_if_distributed for immediate output
- Add train_debug function with rank and local_rank info for training step debugging
- Add debug logs for forward_backward and clip_grad_and_step in first 2 steps
- Enhance fsdp2_load_full_state_dict patching with detailed per-parameter state loading traces

- Rename `train_debug` to `debug_print` for clarity
- Add timestamp to debug messages for better tracing
- Support writing debug logs to file via `TWINKLE_DEBUG_DIR` env var
- Improve tensor debug info in FSDP2 state loading
- Remove redundant accelerate FSDP2 delegation path

Co-authored-by: Copilot <copilot@github.com>

- Fix accelerate FSDP2 patch to use `distribute_tensor` for CUDA devices
- Refactor native FSDP to handle rank0-only loading and broadcast
- Add EP expert shard specs collection and rank mapping for state dict broadcast
- Fix non-persistent buffer handling to use broadcast instead of restore

…elism

- Fix indentation error in `base.py` for `device_id` assignment
- Add `_native_fsdp_debug` utility function for FSDP debugging
- Implement pre-EP full state dict capture to avoid redundant state_dict calls
- Add debug logging for expert scatter operations in FSDP

Add `_ep_trace` function and debug logging to MoE expert parallelism operations for better debugging and performance analysis. The trace logs all-to-all and all-gather operations with tensor shapes and split sizes, controlled by `TWINKLE_EP_DEBUG` or `TWINKLE_FSDP_DEBUG` environment variables.

Add a `tag` parameter to `_AllToAll` and `debug_tag` to token permutation functions to enhance traceability of all-to-all operations in both forward and backward passes. This enables better debugging by identifying specific all-to-all calls in trace logs.

- Return `permuted_tokens` instead of `torch.empty_like` in EP forward to maintain backward all-to-all path
- Add `_apply_gate` support for custom gating in MoE expert computation
- Broadcast source metadata (shapes/dtypes) from rank 0 in FSDP to ensure correct dtype for EP expert tensors
- Validate source metadata consistency to prevent silent dtype/shape mismatches during state dict broadcast

Remove the `_ep_trace` debug logging function and all its calls from expert parallelism utilities. This cleanup eliminates verbose debug output controlled by `TWINKLE_EP_DEBUG` environment variable, simplifying the code and removing unnecessary I/O overhead in production. Also removes the `tag` parameter from `_AllToAll` and `all_to_all` functions that was only used for debug tracing.

Move expert parameter tensors to GPU lazily during chunk scattering instead of loading all at once, preventing OOM when handling large expert parameters.

…ility